import os
import gradio as gr
from dialog import DialogModel
from rag import FileManager, RAGManager
import threading
import asyncio
import shutil
from emotion_analysis import EmoAnal  # 导入心情分析类
import json
from joke_generator import JokeGenerator

prompt_template_file = "data/prompt-template.json"

class FileManagerSystem:
    def __init__(self):
        self.file_manager = FileManager()
        self.rag_manager = RAGManager()
        self.dialog_model = DialogModel()
        self.emo_anal = EmoAnal()
        self.file_lock = threading.Lock()
        self.joke_gen = JokeGenerator()

    def predict(self, text, chat_history, use_rag, selected_file=None):
        """处理用户输入并生成回复"""
        updated_files = self.update_file_list()
        text_with_rag = text
        do_sample = True
        temperature = 0.7

        if use_rag and selected_file:
            file_path = os.path.join(self.file_manager.upload_dir, selected_file)
            file_content = self.file_manager.read_file(file_path)

            if file_content and "出错" not in file_content:
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                text_with_rag = loop.run_until_complete(self.async_rag_process(text, file_content))
                do_sample = False
                temperature = 0.0
            else:
                chat_history.append((text, "Error reading selected file."))
                return "", chat_history

        chat_history = self.dialog_model.predict(text_with_rag, chat_history, do_sample=do_sample,
                                                 temperature=temperature, max_length=2500)

        # 检查用户心情状态
        status = self.emo_anal.check_keywords(text)
        if "😢" in status or "😐" in status:
            joke = self.joke_gen.get_random_joke()
            # 获取最新的回复
            last_question, last_answer = chat_history[-1]
            # 在回复末尾添加笑话
            new_answer = f"{last_answer}\n\n🤖 每日joke，心晴智伴陪您每天笑一笑🌞：\n{joke}"
            chat_history[-1] = (last_question, new_answer)

        return "", chat_history

    async def async_rag_process(self, text, file_content):
        """异步处理RAG流程"""
        chunks = self.rag_manager.chunk_text(file_content)
        await asyncio.to_thread(self.rag_manager.generate_embeddings)
        retrieved = await asyncio.to_thread(self.rag_manager.retrieve, text)
        return f"{text}\n\n参考内容：\n" + "\n".join(retrieved[:3])

    def update_file_list(self):
        """更新文件列表"""
        files = []
        try:
            if os.path.exists(self.file_manager.upload_dir):
                for f in os.listdir(self.file_manager.upload_dir):
                    path = os.path.join(self.file_manager.upload_dir, f)
                    if os.path.isfile(path):
                        size = f"{os.path.getsize(path) / 1024:.1f}KB"
                        files.append([f, size])
                print(f"File list updated: Found {len(files)} files")
        except Exception as e:
            print(f"Error updating file list: {str(e)}")
        return files

    def on_file_select(self, evt: gr.SelectData):
        """处理文件选择事件"""
        try:
            filename = evt.value
            file_path = os.path.join(self.file_manager.upload_dir, filename)
            return self.file_manager.read_file(file_path), filename
        except Exception as e:
            return f"Preview error: {e}", None

    def delete_selected_file(self, selected_filename):
        """删除选中的文件"""
        if selected_filename:
            file_path = os.path.join(self.file_manager.upload_dir, selected_filename)
            if os.path.exists(file_path):
                try:
                    os.remove(file_path)
                    print(f"File deleted: {file_path}")
                except Exception as e:
                    print(f"Error deleting file: {e}")
        return self.update_file_list(), None

    def after_upload(self, files):
        """处理文件上传后事件"""
        with self.file_lock:
            for file in files:
                temp_path = file.name
                self.file_manager.save_uploaded_file(temp_path)
        return self.update_file_list()

    def open_preview(self, file_list, selected_file):
        """打开文件预览"""
        if selected_file:
            file_path = os.path.join(self.file_manager.upload_dir, selected_file)
            if os.path.exists(file_path):
                with open(file_path, "r", encoding="utf-8") as f:
                    content = f.read()
                return gr.update(visible=True), gr.update(value=content)
            else:
                return gr.update(visible=True), gr.update(value="文件不存在")
        else:
            return gr.update(visible=True), gr.update(value="请选择一个文件进行预览")

    @staticmethod
    def close_preview():
        """关闭文件预览"""
        return gr.update(visible=False), gr.update(value=None)

    def check_emoji(self, text):
        """检查表情符号"""
        return self.emo_anal.check_keywords(text)

    @staticmethod
    def load_first_sentence():
        """加载第一条句子"""
        firstSentence = ""
        if os.path.exists(prompt_template_file):
            with open(prompt_template_file, 'r', encoding='utf-8') as file:
                data = json.load(file)
                if data and isinstance(data, list) and len(data) > 0:
                    firstSentence = data[0].strip()
        return firstSentence


# 初始化核心系统
file_manager_system = FileManagerSystem()

custom_css = """
/* 全局样式 */
body {
    font-family: 'Comic Sans MS', 'Chalkboard SE', 'Marker Felt', sans-serif;
    margin: 0;
    padding: 0;
    background-color: #fff8e1;
}

/* 浅色模式样式 */
@media (prefers-color-scheme: light) {
    body {
        background-color: #fff8e1;
        color: #333;
    }

    #chatbot, #preview-textbox, #attachment-management, #qa-area {
        background-color: #fff3e0;
        box-shadow: 0 4px 15px rgba(0,0,0,0.1);
        border: 1px solid #ffe0b2;
    }

    #delete-btn {
        background-color: #ffcc80;
    }

    #delete-btn:hover {
        background-color: #ffb74d;
    }

    .open-preview-btn {
        background-color: #ffcc80;
    }

    .open-preview-btn:hover {
        background-color: #ffb74d;
    }

    .close-preview-btn {
        background-color: #ffcc80;
    }

    .close-preview-btn:hover {
        background-color: #ffb74d;
    }

    .gr-button {
        background-color: #ffb74d;
        color: white;
    }

    .gr-button:hover {
        background-color: #ffa726;
    }

    h1 {
        color: #e65100;
        text-align: center;
    }

    #emoji-container {
        background-color: rgba(255, 243, 224, 0.7);
        border: 1px solid #ffe0b2;
    }
}

/* 主聊天窗口样式 */
#chatbot { 
    min-height: 400px;
    max-height: 70vh;
    overflow-y: auto !important;
    border-radius: 12px;
    padding: 15px;
}

/* 文档预览区域样式 */
#preview-textbox {
    height: 400px !important;
    contain: strict !important;
    border-radius: 12px;
    margin-top: 15px;
    padding: 15px;
}

/* 输入框布局优化 */
.centered-row {
    align-items: stretch !important;
    gap: 12px;
    margin-top: 20px;
}

/* 现代滚动条样式 */
::-webkit-scrollbar {
    width: 10px;
    height: 10px;
    background-color: rgba(0,0,0,0.05);
}

::-webkit-scrollbar-thumb {
    background-color: rgba(0,0,0,0.2);
    border-radius: 5px;
    border: 2px solid transparent;
    background-clip: content-box;
}

::-webkit-scrollbar-thumb:hover {
    background-color: rgba(0,0,0,0.3);
}

/* 左侧附件管理区域样式 */
#attachment-management {
    padding-right: 20px;
    border-radius: 12px;
    padding: 20px;
}

/* 右侧问答区域样式 */
#qa-area {
    padding-left: 20px;
    border-radius: 12px;
    padding: 20px;
}

/* 文件列表样式 */
#file-list {
    border: 1px solid #ffe0b2;
    border-radius: 8px;
    margin-top: 15px;
    padding: 10px;
}

/* 删除按钮样式 */
#delete-btn {
    margin-top: 15px;
    width: 100%;
    color: white;
    border: none;
    border-radius: 8px;
    padding: 10px;
    cursor: pointer;
    font-weight: 500;
}

#delete-btn:hover {
    opacity: 0.9;
}

/* 预览按钮样式 */
.open-preview-btn {
    margin-top: 15px;
    width: 100%;
    color: white;
    border: none;
    border-radius: 8px;
    padding: 10px;
    cursor: pointer;
    font-weight: 500;
}

.open-preview-btn:hover {
    opacity: 0.9;
}

.close-preview-btn {
    margin-top: 15px;
    width: 100%;
    color: white;
    border: none;
    border-radius: 8px;
    padding: 10px;
    cursor: pointer;
    font-weight: 500;
}

.close-preview-btn:hover {
    opacity: 0.9;
}

/* 聊天输入区域样式 */
#chat-input {
    margin-top: 25px;
}

/* 标题样式 */
h1 {
    text-align: center;
    font-size: 36px;
    font-weight: bold;
    margin-bottom: 30px;
}

/* 按钮样式 */
.gr-button {
    border: none;
    border-radius: 8px;
    padding: 10px 15px;
    font-weight: 500;
    cursor: pointer;
    transition: all 0.3s ease;
}

/* 文本框样式 */
.gr-textbox {
    border: 1px solid #ffe0b2;
    border-radius: 8px;
    padding: 10px;
    font-family: 'Comic Sans MS', 'Chalkboard SE', 'Marker Felt', sans-serif;
}

/* 检查框样式 */
.gr-checkbox {
    margin-top: 15px;
}

/* HTML进度条样式 */
#progress {
    text-align: center;
    margin-top: 15px;
}

#emoji-container {
    height: 50px;
    display: flex;
    justify-content: center;
    align-items: center;
    background-color: rgba(255, 243, 224, 0.7);
    border-radius: 8px;
    margin-bottom: 10px;
    box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);
    font-size: 16px;
    padding: 0 10px;
}
"""

with (gr.Blocks(
        css=custom_css,
        theme=gr.themes.Default(
            font=[gr.themes.GoogleFont("Comic Sans MS"),
                  "Chalkboard SE", "Marker Felt", "sans-serif"],
            primary_hue="orange"
        )
) as userDialogueInterface):
    gr.Markdown("# 🌈 心晴智伴 Q&A系统 🌈")
    selected_file = gr.State()
    preview_visible = gr.State(False)

    userDlgInf = userDialogueInterface

    with gr.Row():
        with gr.Column(scale=3, elem_id="attachment-management"):
            gr.Markdown("## 📁 附件管理")
            file_upload = gr.File(
                label="📤 上传附件",
                file_types=[".txt", ".pdf", ".docx", ".md"],
                file_count="multiple"
            )
            file_list = gr.Dataframe(
                headers=["📄 附件名", "📏 附件大小"],
                datatype=["str", "str"],
                interactive=True,
                elem_id="file-list",
                height=300
            )
            delete_btn = gr.Button(
                "🗑️ 删除所选附件",
                elem_id="delete-btn",
                variant="secondary"
            )

            open_preview_btn = gr.Button("📖 打开文档预览", elem_classes="open-preview-btn")
            close_preview_btn = gr.Button(
                "❌ 关闭文档预览",
                elem_classes="close-preview-btn",
                visible=False
            )
            preview_column = gr.Column(visible=False)
            with preview_column:
                file_preview = gr.Textbox(
                    label="📄 文档内容",
                    interactive=False,
                    elem_id="preview-textbox"
                )

        with gr.Column(scale=7, elem_id="qa-area"):
            gr.Markdown("## ❓ 问答区")
            use_rag = gr.Checkbox(label="📚 启用文档Q&A", value=True)
#            progress = gr.HTML("<div id='progress' style='color:gray'>Ready</div>")

            with gr.Column():
                emoji_container = gr.Column(scale=1, elem_id="emoji-container")
                with emoji_container:
                    emoji_display = gr.HTML("当前状态 : 😐")
                chatbot = gr.Chatbot(label="💬 对话历史", elem_id="chatbot")

            with gr.Row(elem_classes="centered-row", elem_id="chat-input"):
                msg = gr.Textbox(
                    label="输入消息",
                    show_label=False,
                    placeholder="✏️ 输入消息并按回车或发送",
                    container=False,
                    scale=8,
                    lines=1
                )
                with gr.Column(scale=1, min_width=150):
                    with gr.Row():
                        submit_btn = gr.Button(
                            "▶️",
                            elem_classes="icon-btn",
                            value="发送"
                        )
                        clear_btn = gr.Button(
                            "✖️",
                            elem_classes="icon-btn",
                            value="清除历史"
                        )

    # 文件处理相关事件
    file_upload.upload(
        file_manager_system.after_upload,
        inputs=file_upload,
        outputs=[file_list]
    )
    file_list.select(
        file_manager_system.on_file_select,
        outputs=[file_preview, selected_file]
    )
    delete_btn.click(
        file_manager_system.delete_selected_file,
        inputs=[selected_file],
        outputs=[file_list, selected_file]
    )

    open_preview_btn.click(
        file_manager_system.open_preview,
        inputs=[file_list, selected_file],
        outputs=[preview_column, file_preview]
    )

    close_preview_btn.click(
        file_manager_system.close_preview,
        outputs=[preview_column, file_preview]
    )

    # 对话处理相关事件
    submit_btn.click(
        file_manager_system.predict,
        inputs=[msg, chatbot, use_rag, selected_file],
        outputs=[msg, chatbot],
        api_name=False,
        queue=True
    )
    msg.submit(
        file_manager_system.predict,
        inputs=[msg, chatbot, use_rag, selected_file],
        outputs=[msg, chatbot],
        api_name=False,
        queue=True
    )

    clear_btn.click(
        lambda: ([], None),
        outputs=[chatbot, selected_file]
    )
    userDlgInf.load(file_manager_system.update_file_list, outputs=[file_list])

    # 表情符号检查事件
    submit_btn.click(
        file_manager_system.check_emoji,
        inputs=msg,
        outputs=emoji_display
    )
    msg.submit(
        file_manager_system.check_emoji,
        inputs=msg,
        outputs=emoji_display
    )

    # 自动发送第一句
    first_sentence = file_manager_system.load_first_sentence()
    if first_sentence:
        msg.value = first_sentence

    # 添加鼠标移动事件监听器
    userDlgInf.load(lambda: """
    <script>
    let hasMoved = false;
    function onMouseMove() {
        if (!hasMoved) {
            hasMoved = true;
            document.querySelector('button[value="发送"]').click();
        }
    }
    document.addEventListener('mousemove', onMouseMove);
    </script>
    """)

# 清理缓存
cache_path = os.path.expanduser("~/.cache/torch_extensions")
if os.path.exists(cache_path):
    shutil.rmtree(cache_path)

if __name__ == "__main__":
    print(f"上传目录: {file_manager_system.file_manager.upload_dir}")
    print("启动 Gradio 应用...")
    userDlgInf.launch(share=True, max_threads=10)